#include <gtest/gtest.h>
#include <chrono>
#include "viml/parser.h"
#include "viml/json_encoder.h"

namespace viml {

class JSONASTTest : public ::testing::Test {
protected:
    void SetUp() override {
        // 准备测试数据
    }
    
    void TearDown() override {
        // 清理
    }
    
    // 辅助函数：解析JSON字符串并验证结构
    bool containsKey(const std::string& json, const std::string& key) {
        return json.find("\"" + key + "\":") != std::string::npos;
    }
    
    bool containsValue(const std::string& json, const std::string& value) {
        return json.find(value) != std::string::npos;
    }
};

// 测试基本的JSON输出功能
TEST_F(JSONASTTest, BasicNumberLiteral) {
    std::string source = "42";
    viml::Parser parser(source);
    auto program = parser.parse();
    
    std::string json = program->to_json();
    
    // 验证JSON包含期望的内容
    EXPECT_TRUE(containsKey(json, "type"));
    EXPECT_TRUE(containsValue(json, "\"Program\""));
    EXPECT_TRUE(containsKey(json, "statements"));
    EXPECT_TRUE(containsValue(json, "\"NumberLiteral\""));
    EXPECT_TRUE(containsValue(json, "42"));
}

// 测试字符串字面量的JSON输出
TEST_F(JSONASTTest, StringLiteral) {
    std::string source = "\"hello world\"";
    viml::Parser parser(source);
    auto program = parser.parse();
    
    std::string json = program->to_json();
    
    EXPECT_TRUE(containsValue(json, "\"StringLiteral\""));
    EXPECT_TRUE(containsValue(json, "\"hello world\""));
}

// 测试标识符的JSON输出
TEST_F(JSONASTTest, Identifier) {
    std::string source = "variable_name";
    viml::Parser parser(source);
    auto program = parser.parse();
    
    std::string json = program->to_json();
    
    EXPECT_TRUE(containsValue(json, "\"Identifier\""));
    EXPECT_TRUE(containsKey(json, "name"));
    EXPECT_TRUE(containsValue(json, "\"variable_name\""));
}

// 测试二元表达式的JSON输出
TEST_F(JSONASTTest, BinaryExpression) {
    std::string source = "1 + 2";
    viml::Parser parser(source);
    auto program = parser.parse();
    
    std::string json = program->to_json();
    
    EXPECT_TRUE(containsValue(json, "\"BinaryExpression\""));
    EXPECT_TRUE(containsKey(json, "operator"));
    EXPECT_TRUE(containsValue(json, "\"+\""));
    EXPECT_TRUE(containsKey(json, "left"));
    EXPECT_TRUE(containsKey(json, "right"));
    EXPECT_TRUE(containsValue(json, "\"NumberLiteral\""));
}

// 测试复杂二元表达式的JSON输出
TEST_F(JSONASTTest, ComplexBinaryExpression) {
    std::string source = "1 + 2 * 3";
    viml::Parser parser(source);
    auto program = parser.parse();
    
    std::string json = program->to_json();
    
    // 应该包含嵌套的二元表达式
    EXPECT_TRUE(containsValue(json, "\"BinaryExpression\""));
    EXPECT_TRUE(containsValue(json, "\"+\""));
    EXPECT_TRUE(containsValue(json, "\"*\""));
    
    // 验证操作数的结构
    size_t count = 0;
    size_t pos = json.find("\"NumberLiteral\"");
    while (pos != std::string::npos) {
        count++;
        pos = json.find("\"NumberLiteral\"", pos + 1);
    }
    EXPECT_EQ(count, 3); // 应该有3个数字字面量
}

// 测试Let语句的JSON输出
TEST_F(JSONASTTest, LetStatement) {
    std::string source = "let x = 42";
    viml::Parser parser(source);
    auto program = parser.parse();
    
    std::string json = program->to_json();
    
    EXPECT_TRUE(containsValue(json, "\"LetStatement\""));
    EXPECT_TRUE(containsKey(json, "name"));
    EXPECT_TRUE(containsValue(json, "\"x\""));
    EXPECT_TRUE(containsKey(json, "value"));
    EXPECT_TRUE(containsValue(json, "\"NumberLiteral\""));
}

// 测试Echo语句的JSON输出
TEST_F(JSONASTTest, EchoStatement) {
    std::string source = "echo \"hello\"";
    viml::Parser parser(source);
    auto program = parser.parse();
    
    std::string json = program->to_json();
    
    EXPECT_TRUE(containsValue(json, "\"EchoStatement\""));
    EXPECT_TRUE(containsKey(json, "expression"));
    EXPECT_TRUE(containsValue(json, "\"StringLiteral\""));
    EXPECT_TRUE(containsValue(json, "\"hello\""));
}

// 测试函数调用的JSON输出
TEST_F(JSONASTTest, CallExpression) {
    std::string source = "func(1, 2)";
    viml::Parser parser(source);
    auto program = parser.parse();
    
    std::string json = program->to_json();
    
    EXPECT_TRUE(containsValue(json, "\"CallExpression\""));
    EXPECT_TRUE(containsKey(json, "callee"));
    EXPECT_TRUE(containsKey(json, "arguments"));
    EXPECT_TRUE(containsValue(json, "\"Identifier\""));
    EXPECT_TRUE(containsValue(json, "\"func\""));
}

// 测试列表字面量的JSON输出
TEST_F(JSONASTTest, ListLiteral) {
    std::string source = "[1, 2, 3]";
    viml::Parser parser(source);
    auto program = parser.parse();
    
    std::string json = program->to_json();
    
    EXPECT_TRUE(containsValue(json, "\"ListLiteral\""));
    EXPECT_TRUE(containsKey(json, "elements"));
    
    // 验证包含3个元素
    size_t count = 0;
    size_t pos = json.find("\"NumberLiteral\"");
    while (pos != std::string::npos) {
        count++;
        pos = json.find("\"NumberLiteral\"", pos + 1);
    }
    EXPECT_EQ(count, 3);
}

// 测试字典字面量的JSON输出
TEST_F(JSONASTTest, DictLiteral) {
    std::string source = "{\"key\": \"value\"}";
    viml::Parser parser(source);
    auto program = parser.parse();
    
    std::string json = program->to_json();
    
    EXPECT_TRUE(containsValue(json, "\"DictLiteral\""));
    EXPECT_TRUE(containsKey(json, "pairs"));
    EXPECT_TRUE(containsKey(json, "key"));
    EXPECT_TRUE(containsKey(json, "value"));
    EXPECT_TRUE(containsValue(json, "\"StringLiteral\""));
}

// 测试函数定义的JSON输出
TEST_F(JSONASTTest, FunctionStatement) {
    // 使用简单的单行函数进行测试
    std::string source = "function add(a, b)\nreturn a + b\nendfunction";
    viml::Parser parser(source);
    auto program = parser.parse();
    
    std::string json = program->to_json();
    
    // 打印JSON以便调试
    std::cout << "Function JSON: " << json << std::endl;
    
    EXPECT_TRUE(containsValue(json, "\"Program\""));
    EXPECT_TRUE(containsKey(json, "statements"));
}

// 测试条件语句的JSON输出
TEST_F(JSONASTTest, IfStatement) {
    std::string source = "if x > 0\necho \"positive\"\nelse\necho \"negative\"\nendif";
    viml::Parser parser(source);
    auto program = parser.parse();
    
    std::string json = program->to_json();
    
    EXPECT_TRUE(containsValue(json, "\"IfStatement\""));
    EXPECT_TRUE(containsKey(json, "condition"));
    EXPECT_TRUE(containsKey(json, "then_body"));
    EXPECT_TRUE(containsKey(json, "else_body"));
    EXPECT_TRUE(containsValue(json, "\"BinaryExpression\""));
    EXPECT_TRUE(containsValue(json, "\">\""));
}

// 测试循环语句的JSON输出
TEST_F(JSONASTTest, WhileStatement) {
    std::string source = "while i\necho 1\nendwhile";
    viml::Parser parser(source);
    auto program = parser.parse();
    
    std::string json = program->to_json();
    
    std::cout << "While JSON: " << json << std::endl;
    
    EXPECT_TRUE(containsValue(json, "\"Program\""));
    EXPECT_TRUE(containsKey(json, "statements"));
}

// 测试JSON编码器的字符串转义功能
TEST_F(JSONASTTest, JSONEncoderStringEscape) {
    std::string test_str = "Hello\\n\\t\\\"World\\\"";
    std::string escaped = JSONEncoder::escape_string(test_str);
    
    EXPECT_TRUE(escaped.find("\\\\n") != std::string::npos);
    EXPECT_TRUE(escaped.find("\\\\t") != std::string::npos);
    EXPECT_TRUE(escaped.find("\\\\\\\"") != std::string::npos);
}

// 测试JSON编码器的对象创建功能
TEST_F(JSONASTTest, JSONEncoderObject) {
    std::vector<std::pair<std::string, std::string>> fields = {
        {"type", JSONEncoder::string_value("test")},
        {"value", JSONEncoder::number_value(42)},
        {"flag", JSONEncoder::bool_value(true)}
    };
    
    std::string json = JSONEncoder::object(fields);
    
    EXPECT_TRUE(json.find("{") != std::string::npos);
    EXPECT_TRUE(json.find("}") != std::string::npos);
    EXPECT_TRUE(json.find("\"type\":\"test\"") != std::string::npos);
    EXPECT_TRUE(json.find("\"value\":42") != std::string::npos);
    EXPECT_TRUE(json.find("\"flag\":true") != std::string::npos);
}

// 测试JSON编码器的数组创建功能
TEST_F(JSONASTTest, JSONEncoderArray) {
    std::vector<std::string> elements = {
        JSONEncoder::string_value("first"),
        JSONEncoder::string_value("second"),
        JSONEncoder::number_value(3)
    };
    
    std::string json = JSONEncoder::array(elements);
    
    EXPECT_TRUE(json.find("[") != std::string::npos);
    EXPECT_TRUE(json.find("]") != std::string::npos);
    EXPECT_TRUE(json.find("\"first\"") != std::string::npos);
    EXPECT_TRUE(json.find("\"second\"") != std::string::npos);
    EXPECT_TRUE(json.find("3") != std::string::npos);
}

// 测试复杂嵌套结构的JSON输出
TEST_F(JSONASTTest, ComplexNestedStructure) {
    std::string source = "let data = {\"numbers\": [1, 2, 3], \"name\": \"test\"}";
    viml::Parser parser(source);
    auto program = parser.parse();
    
    std::string json = program->to_json();
    
    EXPECT_TRUE(containsValue(json, "\"LetStatement\""));
    EXPECT_TRUE(containsValue(json, "\"DictLiteral\""));
    EXPECT_TRUE(containsValue(json, "\"ListLiteral\""));
    EXPECT_TRUE(containsValue(json, "\"StringLiteral\""));
    EXPECT_TRUE(containsValue(json, "\"NumberLiteral\""));
    
    // 验证嵌套结构完整性
    EXPECT_TRUE(containsKey(json, "pairs"));
    EXPECT_TRUE(containsKey(json, "elements"));
}

// 性能测试 - 确保JSON序列化不会显著影响性能
TEST_F(JSONASTTest, PerformanceTest) {
    // 创建一个相对复杂的VimL代码
    std::string source = R"(
        function fibonacci(n)
            if n <= 1
                return n
            else
                return fibonacci(n-1) + fibonacci(n-2)
            endif
        endfunction
        
        let result = fibonacci(10)
        echo result
    )";
    
    viml::Parser parser(source);
    auto program = parser.parse();
    
    auto start = std::chrono::high_resolution_clock::now();
    
    std::string json = program->to_json();
    
    auto end = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start);
    
    // 确保JSON序列化时间在合理范围内（小于50毫秒）
    EXPECT_LT(duration.count(), 50);
    EXPECT_FALSE(json.empty());
    
    // 验证生成的JSON包含预期结构
    EXPECT_TRUE(containsValue(json, "\"FunctionStatement\""));
    EXPECT_TRUE(containsValue(json, "\"fibonacci\""));
    EXPECT_TRUE(containsValue(json, "\"IfStatement\""));
    EXPECT_TRUE(containsValue(json, "\"CallExpression\""));
}

} // namespace viml